# -*- coding: utf-8 -*-
"""
******* 文档说明 ******
数据处理部分
训练数据通过类以迭代器方式
验证数据直接返回对应列表

传入训练、测试数据 csv 板式
第一列为图片路径，第二列为标签值。 第一行为列名
TrainImage,GroundTrue
D:/Desktop/Cifar10-Classification\Data\Image\train_data\frog_00001.bmp,frog
D:/Desktop/Cifar10-Classification\Data\Image\train_data\truck_00002.bmp,truck
D:/Desktop/Cifar10-Classification\Data\Image\train_data\truck_00003.bmp,truck
D:/Desktop/Cifar10-Classification\Data\Image\train_data\deer_00004.bmp,deer

# 当前项目: Cifar10-Classification
# 创建时间: 2019/6/22 17:13
# 开发作者: vincent
# 创建平台: PyCharm Community Edition
# 版    本: V1.0
"""
import cv2
import numpy as np
import random
from collections import Counter
# from keras.utils import np_utils


def to_categorical(y, num_classes=None, dtype='float32'):
    """Converts a class vector (integers) to binary class matrix.

    E.g. for use with categorical_crossentropy.

    # Arguments
        y: class vector to be converted into a matrix
            (integers from 0 to num_classes).
        num_classes: total number of classes.
        dtype: The data type expected by the input, as a string
            (`float32`, `float64`, `int32`...)

    # Returns
        A binary matrix representation of the input. The classes axis
        is placed last.

    # Example

    ```python
    # Consider an array of 5 labels out of a set of 3 classes {0, 1, 2}:
    > labels
    array([0, 2, 1, 2, 0])
    # `to_categorical` converts this into a matrix with as many
    # columns as there are classes. The number of rows
    # stays the same.
    > to_categorical(labels)
    array([[ 1.,  0.,  0.],
           [ 0.,  0.,  1.],
           [ 0.,  1.,  0.],
           [ 0.,  0.,  1.],
           [ 1.,  0.,  0.]], dtype=float32)
    ```
    """

    y = np.array(y, dtype='int')
    input_shape = y.shape
    if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:
        input_shape = tuple(input_shape[:-1])
    y = y.ravel()
    if not num_classes:
        num_classes = np.max(y) + 1
    n = y.shape[0]
    categorical = np.zeros((n, num_classes), dtype=dtype)
    categorical[np.arange(n), y] = 1
    output_shape = input_shape + (num_classes,)
    categorical = np.reshape(categorical, output_shape)
    return categorical


# 打印日志，若无日志对象直接打印。 方便无日志模块调试
def f_print(string, level='INFO', logger=None):
    if logger is None:
        print('{}:{}'.format(level, string))
    else:
        if level == 'INFO':
            logger.info(string)
        elif level == 'WARNING':
            logger.warning(string)
        elif level == 'ERROR':
            logger.error(string)
        elif level == 'CRITICAL':
            logger.critical(string)
        else:
            logger.debug(string)


# 验证数据
def val_data(val_data_csv_path, input_image_shape, label_mapping, logger=None):
    """
    :param val_data_csv_path:  验证数据 CSV 列表
    :param input_image_shape:  图片 Resize 大小
    :param label_mapping: 类别字典
    :param logger:  日志对象
    :return:
    """
    # 验证数据列表
    val_img_path = list()
    val_img_label = list()
    for img_info in open(val_data_csv_path, 'r', encoding='utf-8'):
        img_info = img_info.strip().split(',')
        # 训练图片路径
        val_img_path.append(img_info[0])
        # 训练图片标签
        val_img_label.append(img_info[1])

    # 第一行为标签名
    val_img_path = val_img_path[1:]
    val_img_label = val_img_label[1:]

    x_val = np.zeros((len(val_img_path), input_image_shape, input_image_shape, 3), dtype='uint8')
    y_val = np.zeros((len(val_img_path), len(label_mapping)), dtype='uint8')

    label_len = len(label_mapping)
    for i, (img_path, img_label) in enumerate(zip(val_img_path, val_img_label)):
        # 读取图片数据
        img_data = cv2.imdecode(np.fromfile(img_path, dtype='uint8'), -1)
        # 图片大小重新设定
        x_val[i] = cv2.resize(img_data, (input_image_shape, input_image_shape))

        # 读取图片标签
        y_val[i] = to_categorical([label_mapping[img_label]], label_len)

    f_print('ValData Label Count:{}'.format(sorted(Counter(val_img_label).items(),
                                                   key=lambda x: [x[1], x[1]], reverse=True)), logger=logger)
    return x_val, y_val


# 训练数据读取
class TrainData(object):
    def __init__(self, train_data_csv_path, logger=None):
        """
        :param train_data_csv_path:  训练图片路径、标签 CSV 文件。
        第一列为图片路径，第二列为标签值。 第一行为列名
        TrainImage,GroundTrue
        D:/Desktop/Cifar10-Classification\Data\Image\train_data\frog_00001.bmp,frog
        :param logger:   日志对象
        """

        # 读取训练数据列表
        train_img_path = list()
        train_img_label = list()
        for img_info in open(train_data_csv_path, 'r', encoding='utf-8'):
            img_info = img_info.strip().split(',')
            # 训练图片路径
            train_img_path.append(img_info[0])
            # 训练图片标签
            train_img_label.append(img_info[1])

        # 第一行为标签名
        self.train_img_path = train_img_path[1:]
        self.train_img_label = train_img_label[1:]

        # 各类别占比  按类别名称、类别数量升序
        label_count = sorted(Counter(self.train_img_label).items(), key=lambda x: [x[0], x[1]], reverse=False)
        # 类别映射
        self.label_mapping = {label_i[0]: i for i, label_i in enumerate(label_count)}
        # 类别数量
        self.label_len = len(self.label_mapping)
        # 类别索引映射
        self.label_index_mapping = {_value: _key for _key, _value in self.label_mapping.items()}

        f_print('Label Mapping:{}'.format(self.label_index_mapping), logger=logger)
        f_print('TrainData Label Count:{}'.format(label_count), logger=logger)
        
    # 迭代器 生成批次训练数据
    def generators(self, input_image_shape=224, output_image_label=10, batch_size=10):
        """
        :param input_image_shape:   图片转换后大小
        :param output_image_label:  图片分类类别数量
        :param batch_size:  批次大小
        :return:
        """
        # 判断训练数据实际标签数量是否等于模型设定标签数量
        assert self.label_len == output_image_label, \
            'Error: TrainDataLabel [{}] is not equal ModelLabel [{}]!!!'.format(self.label_len, output_image_label)
        train_data_list = list(zip(self.train_img_path, self.train_img_label))

        random.shuffle(train_data_list)

        x_train = np.zeros((batch_size, input_image_shape, input_image_shape, 3), dtype='uint8')
        y_train = np.zeros((batch_size, output_image_label), dtype='uint8')
        i = 0
        for img_path, img_label in train_data_list:
            # 读取图片数据
            img_data = cv2.imdecode(np.fromfile(img_path, dtype='uint8'), -1)
            # 图片大小重新设定
            x_train[i] = cv2.resize(img_data, (input_image_shape, input_image_shape))

            # 读取图片标签
            y_train[i] = to_categorical([self.label_mapping[img_label]], output_image_label)
            i += 1
            if i == batch_size:
                i = 0
                yield x_train, y_train
